import os
#import matplotlib
#matplotlib.use('WXAgg')
import numpy as np
import pylab as plt
import processGPR as gpr


class Viewer:
    
    def __init__(self):
        ''' Set parameters '''
        self.cmin = 122
        self.cmax = 142
        
        self.le = None
        self.ri = None
        self.procdir = os.getcwd()
        self.traceNo = 0   
        self.coordsid = None # id for displaying the distance by coordinate in showcoords()
        
        ''' Set up the figure and axes '''
        self.fig = plt.figure(None, figsize = (14, 8))
        
        self.axprofile = plt.axes([0.075, 0.075, 0.65, 0.825])
        self.axtrace = plt.axes([0.75, 0.075, 0.2, 0.825])        
#        self.vline = self.axprofile.axvline(0, linewidth=2, color='r')
        self.tracexlabel = 'Power [dB]'
        
        ''' Connect events '''
#        self.__editmode() # make sure background removal is not applied !!!
        self.__pickmode()
    
    
    def __editmode(self):
        print '''
        Press 'q' to define first trace in trace-range
        Press 'w' to define last trace in trace-range
        Press 'r' to remove trace-range
        Press 'p' to save profile to 'title'.h5
        Press 'n' to go to next profile
        Press 'a' to view the spectrum of the current trace
        '''
        self.cidgetleft = self.fig.canvas.mpl_connect('key_press_event', self.getleft)
        self.cidgetright = self.fig.canvas.mpl_connect('key_press_event', self.getright)
        self.cidremove = self.fig.canvas.mpl_connect('key_press_event', self.remove)
        self.cidsave = self.fig.canvas.mpl_connect('key_press_event', self.save)
        self.cidnext = self.fig.canvas.mpl_connect('key_press_event', self.next)
        self.cidgettrace = self.fig.canvas.mpl_connect('key_press_event', self.getTrace)
        self.cidplotspec = self.fig.canvas.mpl_connect('key_press_event', self.plotSpec)
        self.cidwiggle = self.fig.canvas.mpl_connect('motion_notify_event', self.wiggle)
        self.cidshowcoords = self.fig.canvas.mpl_connect('motion_notify_event', self.showcoords)
#        self.cidspectrum = self.fig.canvas.mpl_connect('motion_notify_event', self.spectrum)    
    
    
    def __pickmode(self):
        print '''
        Press 'x' to add a pick
        Press 'z' to remove a pick
        Press 'p' to save the picks
        Press 'm' to load picks from a file
        Press 'r' to reset all picks
        Press 'n' to go to next profile
        Press 'a' to view the spectrum of the current trace
        '''
        self.cidaddpick = self.fig.canvas.mpl_connect('key_press_event', self.addpick)
        self.cidundopick = self.fig.canvas.mpl_connect('key_press_event', self.undopick)
        self.cidsavepick = self.fig.canvas.mpl_connect('key_press_event', self.savepick)
        self.cidloadpick = self.fig.canvas.mpl_connect('key_press_event', self.loadpick)
        self.cidresetpick = self.fig.canvas.mpl_connect('key_press_event', self.resetpick)
        self.cidnext = self.fig.canvas.mpl_connect('key_press_event', self.next)
        self.cidgettrace = self.fig.canvas.mpl_connect('key_press_event', self.getTrace)
        self.cidplotspec = self.fig.canvas.mpl_connect('key_press_event', self.plotSpec)
        self.cidwiggle = self.fig.canvas.mpl_connect('motion_notify_event', self.wiggle)
        self.cidshowcoords = self.fig.canvas.mpl_connect('motion_notify_event', self.showcoords)
    
        
    def QtOpenFile(self, dir=os.getcwd(), multi=True):
        import sys
        from PyQt4 import QtGui
        
        extfilter = '''Raw C-band Files (*.gYY);;Processed C-band Files (*.bYY);;Ramac Files (*.rd3);;All Files (*.*)'''
        
        app = QtGui.QApplication(sys.argv)
        dlg = QtGui.QFileDialog
        if multi:
            pathlist = dlg.getOpenFileNames(None, "Open Files", dir, extfilter)
            pathlist.sort()
            filelist = []
            for path in pathlist:
                filelist.append(os.path.basename(str(path)))
            dir = os.path.dirname(str(pathlist[0]))
            self.getList(filelist)
        else:
            pathlist = dlg.getOpenFileName(None, "Open File", dir, extfilter)
            filelist = os.path.basename(str(pathlist))
            dir = os.path.dirname(str(pathlist))
            self.getFile(filelist)
        os.chdir(dir)        
    
    
    def wxOpenFile(self, dir=os.getcwd(), multi=True, filter=0):
        import wx
        
        if filter == 0:
            extfilter = "Processed C-band Files (*.h5)|*.h5|All Files (*.*)|*.*|Raw C-band Files (*.gYY)|*.gYY|Raw C-band Files (*.gYX)|*.gYX|Raw C-band Files (*.gXX)|*.gXX|Raw C-band Files (*.gXY)|*.gXY|Processed C-band Files (*.bYY)|*.bYY|Ramac Files (*.rd3)|*.rd3"
        elif filter == 1:
            extfilter = "Pick files (*.pck)|*.pck|All Files (*.*)|*.*"
        else:
            extfilter = "All Files (*.*)|*.*"
            
        app = wx.PySimpleApp()
        
        if multi:
            dlg = wx.FileDialog(None, message='Select files ...', defaultDir=dir, wildcard=extfilter, style=wx.FD_MULTIPLE)  
            if dlg.ShowModal() == wx.ID_OK:
                pathlist = dlg.GetPaths()
                filelist = []
                for path in pathlist:
                    filelist.append(os.path.basename(str(path)))
                filelist.sort()
                os.chdir(os.path.dirname(str(pathlist[0])))
                self.path = os.path.dirname(str(pathlist[0]))
                self.getList(filelist)
            else:
               filename = False
            dlg.Destroy()
        else:
            dlg = wx.FileDialog(None, message='Select a file ...', defaultDir=dir, wildcard=extfilter, style=wx.FD_OPEN)  
            if dlg.ShowModal() == wx.ID_OK:
                self.path =  str(dlg.GetPath())
                datafile = os.path.basename(str(self.path))
                os.chdir(os.path.dirname(str(self.path)))
                self.getFile(datafile)
            else:
               filename = False
            dlg.Destroy()
    
    
    def getList(self, filelist):
        self.container = filelist
        self.containeriter = iter(filelist)
    
    
    def getFile(self, filename):
        self.container = filename
    
    
    def getData(self, data):
        self.container = None
        self.profile = gpr.CTD()
        ''' Set parameters '''
        self.profile.data = data
        self.profile.samples, self.profile.traces = np.shape(self.profile.data)
        self.profile.t = np.arange(0, self.profile.deltaT*self.profile.samples, self.profile.deltaT)*1e9
        self.profile.r = self.profile.t*self.profile.cmed/2.0
#        self.profile.filename = data.filename
 
 
    def getClass(self, C):
        self.container = None
        self.profile = C
    
    
    def loadData(self):
        if self.container == None:
            file = 'empty'
        elif isinstance(self.container, list):
            file = self.containeriter.next()
        elif isinstance(self.container, str):
            file = self.container
            
        if file == 'empty':
            pass
        elif file[-4:-2] == '.g':
            self.loadRAW(file)        
        elif file[-3:] == '.h5':
            self.loadH5(file)                
        elif file[-4:] == '.rd3':
            self.loadRD3(file)   
        else:
            self.profile = gpr.CTD()
            self.profile.readBin(file, 805)
            self.profile.filename = file
            
        self.picks = Picks(self.profile.traces)
        
        
    def loadH5(self, file):
        self.profile = gpr.CTD()
        self.profile.readH5(file)
#        self.profile.subtractMean(700)
#        self.profile.wscs()
#        self.profile.linearGain(1.05)
##        self.profile.equalizeTraces()
#        self.profile.getPlotRange()        
        self.profile.filename = file[:-3]
        self.cmin = self.profile.cmin
        self.cmax = self.profile.cmax
        ''' Uncomment for data without gain (wscs) '''
#        self.cmin = 0
#        self.cmax = 20000
        ''' Uncomment to increase contrast in the upper layers '''
#        self.cmin = 57
#        self.cmax = 64
        ''' Uncoomment to increase contrast in the deeper layers '''
        self.cmin = 108
        self.cmax = 132   
    
    
    def loadRAW(self, file):
        FD = gpr.CFD()
        FD.readRaw(file)
        self.profile = gpr.CTD()
        self.profile.timeDomain(FD.data)
        print 'Traces: %i' % self.profile.traces
#            self.profile.stackTraces(10)
        self.profile.staticCorrection()
#            self.profile.wscs()
        if len(file[7:-4]) == 1:
            self.profile.filename = os.path.split(os.getcwd())[-1]+'_0'+file[7:-4]
        else:
            self.profile.filename = os.path.split(os.getcwd())[-1]+'_'+file[7:-4]
        self.profile.getPlotRange()
        self.cmax = 100000.0
        self.cmin = 10.0
#            self.cmax = 140.0
#            self.cmin = 120.0
    
    
    def loadRD3(self, file):
        self.profile = gpr.IMP()
        self.profile.loadrd3(file[:-4])
        self.profile.subtractDC()
        self.profile.backgroundRemoval()
        self.profile.stackTraces(10)
        self.cmin = -500 #self.profile.data.min(axis=None)
        self.cmax = 500 #self.profile.data.max(axis=None)
        self.tracexlabel = 'Amplitude'
        self.profile.filename = file[:-4]
    
    
    def draw(self):
        ''' Profile axes '''
        self.axprofile.matshow(self.profile.data, aspect='auto', cmap=plt.cm.Blues,  vmin=self.cmin, vmax=self.cmax,
                              extent=[0, self.profile.traces, self.profile.t[-1], self.profile.t[0]])
        self.axprofile.xaxis.set_ticks_position('bottom')
        self.axprofile.set_title(self.profile.filename)
        self.axprofile.set_xlabel('Trace #')
        self.axprofile.set_ylabel('Time [ns]')
        self.axprofile.grid(True)
        self.axprofile.set_autoscale_on(False) 
        
        ''' Trace axes '''
        self.axtrace.plot(self.profile.data[:, self.traceNo], self.profile.t, color='k') 
        self.axtrace.axis([self.cmin, self.cmax, self.axprofile.axis()[2],  self.axprofile.axis()[3]])
        self.axtrace.set_title('Trace # %i' % self.traceNo)  
        self.axtrace.set_xlabel(self.tracexlabel)
        self.axtrace.set_ylabel('Time [ns]') 
        self.axtrace.yaxis.set_label_position('right')
        self.axtrace.yaxis.tick_right()
        self.vline = self.axprofile.axvline(0, linewidth=2, color='r')
        
        self.fig.canvas.draw()     
        
    
    def showpick(self):
        self.picks.pltid = self.axprofile.scatter(self.picks.trace[:self.picks.tell], self.picks.twt[:self.picks.tell], marker='x', s=70, color='r')
        self.fig.canvas.draw()
        
    def showtrack(self):
        try:
            self.figtrack = plt.figure(None, figsize = (7, 7))
            self.axtrack = plt.axes()
            self.axtrack.plot(self.profile.coords.x, self.profile.coords.y)
            self.axtrack.set_title('PS coordinates\nTrue meridian -71$^{\circ}$')
            self.axtrack.set_xlabel('Easting')
            self.axtrack.set_ylabel('Northing')
        except AttributeError:
            pass
        
        
    def show(self):
#        app = wx.PySimpleApp()
#        app.MainLoop()
        self.loadData()
        self.draw()
        self.showtrack()
        plt.show()   
        
    
    ''' KEY EVENTS '''
    
    def switchmode(self, event):
        pass
    
                    
    def addpick(self, event):
        if event.inaxes:
            if event.key == 'x':
                tx = int(np.round(event.xdata))
                self.picks.trace[self.picks.tell] = tx
                self.picks.twt[self.picks.tell] = event.ydata
                self.picks.x[self.picks.tell] = self.profile.coords.x[tx]
                self.picks.y[self.picks.tell] = self.profile.coords.y[tx]
                print 'Added point: %i %f %f %f' % (self.picks.trace[self.picks.tell], self.picks.twt[self.picks.tell], self.picks.x[self.picks.tell], self.picks.y[self.picks.tell])
                self.picks.tell += 1
                if self.picks.pltid != 0:
                    self.picks.pltid.remove()
                self.showpick()
                
                
    def undopick(self, event):
        if event.key == 'z':
            if self.picks.tell > 1:
                self.picks.tell -= 1
                self.picks.pltid.remove()
                self.showpick()
                print 'Removed last pick!'
            else:
                if self.picks.pltid != 0:
                    self.picks.pltid.remove()
                self.picks = Picks(self.profile.traces)
                print 'Picks reset!'
                

    def resetpick(self, event):       
        if event.key == 'r':  
                self.picks.pltid.remove()   
                self.picks = Picks(self.profile.traces)
                print 'Picks reset!'
            
    def savepick(self, event):
        if event.key == 'p':
            self.picks.save('tmp_pick.pck')
            print 'Picks saved!'
    
    def loadpick(self, event):
        if event.key == 'm':
            self.wxOpenFile(dir=self.path, multi=False, filter=1)
            self.picks.load(self.container)
            self.showpick()
            print 'Picks loaded!'
        
    def getleft(self, event):
        ''' Selects the most left trace in a trace range used in 'remove' '''
        if event.inaxes:
            if event.key == 'q':
                self.le = int(event.xdata)
                print self.le
            
            
    def getright(self, event):
        ''' Selects the most right trace in a trace range used in 'remove' '''
        if event.inaxes:
            if event.key == 'w':
                self.ri = int(event.xdata)
                print self.ri
           

    def remove(self, event):
        ''' Removes the traces selected by 'getleft' and 'getright' '''
        if event.key == 'r':
            self.profile.removeTraces([self.le, self.ri])
            self.le = None
            self.ri = None
            self.axprofile.lines.remove(self.vline)
            self.draw()

            
    def save(self, event):
        ''' Saves the shown profile '''
        if event.key == 'p':
            self.profile.writeH5(self.profile.filename+self.profile.procname)
    
    
    def next(self, event):
        ''' If a list of profiles was selected it shows the next one. '''
        if event.key == 'n':
            self.axprofile.clear()
            self.axtrace.clear()
            self.traceNo = 0
            self.coordsid = None
            self.loadData()
            self.draw()
    
    
    def getTrace(self, event):
        ''' Saves the trace in the wiggle window as 'npy' file '''
        if event.inaxes == self.axprofile.axes:
            if event.key == 't':
                if self.container == None:
                    tracefile = 'trace_'+str(self.traceNo)
                else:
                    tracefile = self.axprofile.get_title()+'_'+str(self.traceNo)
                np.save(tracefile, self.profile.data[:,self.traceNo])
                plt.savefig(tracefile+'.png', dpi=50)
                print 'Trace written to %s %s' % (os.getcwd(), tracefile)
        
    
    def plotSpec(self, event):
        if event.inaxes == self.axprofile.axes:
            if event.key == 'a':
                self.traceNo = int(event.xdata)
                trace = gpr.TRACE(self.profile.data[:,self.traceNo], self.profile.deltaT)
                trace.spectrum()
                plt.show()
                
    
    ''' MOTION EVENTS '''
    
    def wiggle(self, event):
        if event.inaxes == self.axprofile.axes:
            try:
                self.traceNo = int(event.xdata)
                self.axprofile.lines.remove(self.vline)
                self.vline = self.axprofile.axvline(event.xdata, linewidth=2, color='r')
                
                self.axtrace.hold(False)
                self.axtrace.plot(self.profile.data[:,self.traceNo], self.profile.t, color='k') 
#                self.axtrace.plot(self.profile.data[:,self.traceNo], np.arange(self.profile.samples), color='k') 
                self.axtrace.hold(True)  
                self.axtrace.axhline(event.ydata, linewidth=2, color='r') 
                self.axtrace.axis([self.cmin, self.cmax, self.axprofile.axis()[2],  self.axprofile.axis()[3]])
                self.axtrace.set_title('Trace # %i' % self.traceNo)  
                self.axtrace.set_xlabel(self.tracexlabel)
                self.axtrace.set_ylabel('Time [ns]')   
                self.axtrace.yaxis.set_label_position('right')
                self.axtrace.yaxis.tick_right()     
                self.fig.canvas.draw()
            except OverflowError:
                print 'No valid data!'
                
    
    def spectrum(self, event):
        if event.inaxes == self.axprofile.axes:
            try:
                self.traceNo = int(event.xdata)
                self.axprofile.lines.remove(self.vline)
                self.vline = self.axprofile.axvline(event.xdata, linewidth=2, color='r')
                
                
                f = np.fft.fft(self.profile.data[:,self.traceNo])
                amp = np.abs(np.real(f[0:self.profile.samples/2]))
                self.axtrace.hold(False)
                self.axtrace.plot(amp, self.profile.freq, color='k') 
#                self.axtrace.hold(True)  
#                self.axtrace.axhline(event.ydata, linewidth=2, color='r') 
#                self.axtrace.axis([0, 1000, 0, np.abs(np.real(f[-1]))])
                self.axtrace.set_title('Trace # %i' % self.traceNo)  
                self.axtrace.set_xlabel('Amplitude')
                self.axtrace.set_ylabel('Frequency [GHz]')   
                self.axtrace.yaxis.set_label_position('right')
                self.axtrace.yaxis.tick_right()     
                self.fig.canvas.draw()
            except OverflowError:
                print 'No valid data!'

    def showcoords(self, event):
        if event.inaxes == self.axprofile.axes:
            if self.coordsid != None:
                self.coordsid.remove()
            try:
                dist2p = np.add.reduce(self.profile.coords.dist[0:self.traceNo], axis=0)
                self.coordsid = self.axprofile.text(0.5, 1.02, 'Distance from start: %8.2f m' % dist2p,
                                                    horizontalalignment='center',
                                                    verticalalignment='center',
                                                    transform = self.axprofile.transAxes)
            except AttributeError:
                pass


class Picks:    
    def __init__(self, N):
        self.trace = np.empty(N)
        self.twt = np.empty(N)
        self.x = np.empty(N)
        self.y = np.empty(N)
        self.tell = 0 # counts the picks
        self.pltid = 0 # controls if picks are plotted
         
    def save(self, file):
        fid = open(file, 'w')
        for n in range(self.tell):
            fid.write('%i %f %f %f\n' % (self.trace[n], self.twt[n], self.x[n], self.y[n]))
        fid.close()    
        
    def load(self, file):
        fid = open(file, 'r')
        data = fid.readlines()
        fid.close()
        self.tell = len(data)
        for n in range(self.tell):
            self.trace[n], self.twt[n], self.x[n], self.y[n] = data[n].split(' ')
    

if __name__ == '__main__':
    VU = Viewer()
    VU.wxOpenFile('Y:/Data2007_08/CBAND/scambos')
#    VU.getFile('M:/workspace/HUBRA/src/data0.h5')
    VU.show()
    
